/**
 *  Copyright (c) 2024 Ebaina
 *  hieuler u-boot is licensed under Mulan PSL v2.
 *  You can use this software according to the terms and conditions of the Mulan PSL v2. 
 *  You may obtain a copy of Mulan PSL v2 at:
 *           http://license.coscl.org.cn/MulanPSL2 
 *  THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER 
 *  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY OR 
 *  FIT FOR A PARTICULAR PURPOSE.  
 *  See the Mulan PSL v2 for more details.  
 */

#include <common.h>
#include <command.h>
#include <string.h>
#include <env.h>
#include <fs.h>
#include "hex.h"
#include <memalign.h>

#define CMD_CPFILE_HELP "cpfile -s "

#define CMD_MAX_SUB_ARGS 10

/*
cpfile --src-part "src" "dst" 

文件系统中拷贝文件
mmc 1:1 /boot/rootfs.ext4
无文件系统
mmc 1 0x00 [size]

*/

enum HANDLE_TYPE{
    HANDLE_TYPE_ERR,
    HANDLE_TYPE_FLAT,
    HANDLE_TYPE_FS
};

struct rw_handle{
    enum HANDLE_TYPE type;
    struct blk_desc* bdesc;
    char* ifname;
    char* dev_part_str;
    char* filename;
    loff_t offset;
    loff_t size;
    loff_t first;
};

#define startWith(haystack,needle) (strstr(haystack,needle)==haystack)
#define CK_BOOL_FALG(flags,pos) (flags&(0x01<<pos))

static int xhandle(struct rw_handle *handle,char* str)
{
    char *split[CMD_MAX_SUB_ARGS];
    int flag;
    int i,j;

    flag = 0;// 标记引号
    split[0] = str; j = 1;
    for(i=0;str[i];i++){
        if(str[i]==' '&&flag==0){
            str[i]='\0';
            split[j++] = &str[i+1];
        }
        if(str[i]=='\''||str[i]=='\"')
            flag = flag?0:1;
    }

    memset(handle,0,sizeof(struct rw_handle));

    if(split[2][0]=='/'){
        handle->type = HANDLE_TYPE_FS;
        handle->ifname = split[0];
        handle->dev_part_str = split[1];
        handle->filename = split[2];
        handle->offset = 0;
        if (fs_set_blk_dev(handle->ifname,handle->dev_part_str, FS_TYPE_ANY))
            return -1;
        handle->size = 0;
        fs_size(handle->filename,&handle->size);
        handle->first = 0;
    }
    else{
        handle->type = HANDLE_TYPE_FLAT;
        handle->ifname = split[0];
        char*dev = split[1];
        char*part = NULL;
        for(i=0;dev[i]!='\0';i++){
            if(dev[i]==':'){
                part = &dev[i+1];
                dev[i] = '\0';
                break;
            }
        }
        handle->bdesc = blk_get_dev(handle->ifname,atoi_auto(dev));
        handle->offset = 0;
        if(part!=NULL){
            disk_partition_t info;
            part_get_info(handle->bdesc,atoi_auto(part),&info);
            handle->offset += info.start * info.blksz;
        }
        handle->offset += atoi_auto(split[2]);
        handle->size = atoi_auto(split[3]);
        handle->first = handle->offset;
    }

    return 0;
}

static __attribute__((unused)) int show_handle(struct rw_handle *handle)
{
    if(handle->type==HANDLE_TYPE_FS){
        printf("\ttype:HANDLE_TYPE_FS\n");
        printf("\tifname:%s\n",handle->ifname);
        printf("\tdev_part_str:%s\n",handle->dev_part_str);
        printf("\tfilename:%s\n",handle->filename);
        printf("\tsize:0x%llx\n",handle->size);
    }
    else{
        printf("\ttype:HANDLE_TYPE_FLAT\n");
        printf("\tifname:%s\n",handle->ifname);
        printf("\tbdesc:%p\n",handle->bdesc);
        printf("\toffset:0x%llx\n",handle->offset);
        printf("\tsize:0x%llx\n",handle->size);
    }
    return 0;
}

static loff_t rw_handle_read(struct rw_handle *handle,void* buffer,size_t count)
{
    loff_t actread = 0;
    loff_t tmp_count = 0;

    char* _buffer = (char*)buffer;

    if(count==0)
        return 0;
    
    if(handle->type==HANDLE_TYPE_FS){
        if (fs_set_blk_dev(handle->ifname,handle->dev_part_str, FS_TYPE_ANY))
		    return 0;
        fs_read(handle->filename,(ulong)_buffer,handle->offset,count,&actread);
        handle->offset += actread;
    }
    else if(handle->type==HANDLE_TYPE_FLAT){
        if(handle->offset >= handle->first + handle->size)
            return 0;
        if(handle->offset + count >= handle->first + handle->size)
            count = handle->first + handle->size - handle->offset;
        actread = count;

        if(handle->offset % handle->bdesc->blksz != 0){
            char blk_tmp[handle->bdesc->blksz];
            blk_dread(
                handle->bdesc,
                handle->offset / handle->bdesc->blksz,
                1,
                blk_tmp);
            
            tmp_count = handle->bdesc->blksz - (handle->offset % handle->bdesc->blksz);

            memcpy(_buffer,
                blk_tmp+(handle->offset % handle->bdesc->blksz),
                tmp_count);
            
            _buffer += tmp_count;
            count -= tmp_count;
            handle->offset += tmp_count;
        }

        blk_dread(
            handle->bdesc,
            handle->offset / handle->bdesc->blksz,
            count / handle->bdesc->blksz,
            _buffer);
        
        tmp_count = (count / handle->bdesc->blksz) * handle->bdesc->blksz;

        _buffer += tmp_count;
        count -= tmp_count;
        handle->offset += tmp_count;

        if(count != 0){
            char blk_tmp[handle->bdesc->blksz];
            blk_dread(
                handle->bdesc,
                handle->offset / handle->bdesc->blksz,
                1,
                blk_tmp);
            
            memcpy(_buffer,blk_tmp,count);
        }

        handle->offset += count;
    }
    return actread;
}
static loff_t rw_handle_write(struct rw_handle *handle,void* buffer,size_t count)
{
    loff_t actwrite = 0;
    loff_t tmp_count = 0;

    char* _buffer = (char*)buffer;

    if(count==0)
        return 0;

    if(handle->type==HANDLE_TYPE_FS){
        if (fs_set_blk_dev(handle->ifname,handle->dev_part_str, FS_TYPE_ANY))
		    return 0;
        fs_write(handle->filename,(ulong)buffer,handle->offset,count,&actwrite);
        handle->offset += actwrite;
    }
    else if(handle->type==HANDLE_TYPE_FLAT){
        // if(handle->offset >= handle->first + handle->size)
        //     return 0;
        // if(handle->offset + count >= handle->first + handle->size)
        //     count = handle->first + handle->size - handle->offset;
        actwrite = count;

        if(handle->offset % handle->bdesc->blksz != 0){
            char blk_tmp[handle->bdesc->blksz];
            blk_dread(
                handle->bdesc,
                handle->offset / handle->bdesc->blksz,
                1,
                blk_tmp);
            
            tmp_count = handle->bdesc->blksz - (handle->offset % handle->bdesc->blksz);

            memcpy(blk_tmp+(handle->offset % handle->bdesc->blksz),
                _buffer,
                tmp_count);
            
            blk_dwrite(
                handle->bdesc,
                handle->offset / handle->bdesc->blksz,
                1,
                blk_tmp);
            
            _buffer += tmp_count;
            count -= tmp_count;
            handle->offset += tmp_count;
        }

        blk_dwrite(
            handle->bdesc,
            handle->offset / handle->bdesc->blksz,
            count / handle->bdesc->blksz,
            _buffer);
        
        tmp_count = (count / handle->bdesc->blksz) * handle->bdesc->blksz;
        _buffer += tmp_count;
        count -= tmp_count;
        handle->offset += tmp_count;

        if(count != 0){
            char blk_tmp[handle->bdesc->blksz];
            blk_dread(
                handle->bdesc,
                handle->offset / handle->bdesc->blksz,
                1,
                blk_tmp);
            
            memcpy(blk_tmp,_buffer,count);

            blk_dwrite(
                handle->bdesc,
                handle->offset / handle->bdesc->blksz,
                1,
                blk_tmp);
        }

        handle->offset += count;
    }
    return actwrite;
}

void printProgressBar(loff_t progress, loff_t total) {
    const int barWidth = 50;
    int progressChars = (int)((double)progress / total * barWidth);

    printf("[");
    for (int i = 0; i < barWidth; ++i) {
        if ((i+1) < progressChars)
            printf("=");
        else if(i < progressChars)
            printf(">");
        else
            printf(" ");
    }
    printf("] %3d%%\r", (int)((double)progress / total * 100));
}

#define CP_TRUE 0
#define CP_FALSE 1
enum ARGS {
    ARGS_SELF,
    ARGS_SRC,
    ARGS_DST,
    ARGS_OPT
};

int do_cpfile(cmd_tbl_t *cmdtp, int flag, int argc, char * const args[])
{
    // 源与目标
    char buf_src[128];
    char buf_dst[128];
    struct rw_handle hsrc;
    struct rw_handle hdst;

    // 缓冲区
    char* buffer = NULL;
    loff_t buffer_size = 0;

    // 其他
    int i;

    // 缓冲区解析
    if(env_get("buffer_addr")!=NULL)
        buffer = (char*)atoi_auto(env_get("buffer_addr"));
    if(env_get("buffer_size")!=NULL)
        buffer_size = atoi_auto(env_get("buffer_size"));
    
    if(buffer_size==0)
        return CP_FALSE;

    // 源与目标解析
    strcpy(buf_dst,args[ARGS_DST]);
    xhandle(&hdst,buf_dst);

    loff_t actread = 0;
    loff_t actwrite = 0;

    loff_t total = 0;
    loff_t count = 0;

    // 0: 无后缀
    // 1: 禁止无意义的0 rootfs.ext40
    // 2: 以两位数字结尾 rootfs.ext400
    int part_suffix_type = -1;
    const char* part_type[] = {"%s","%s%d","%s%02d"}; 
    // 检测后缀类型
    for (i = 0; i < 3; i++ ) {
        snprintf(buf_src,128,part_type[i],args[ARGS_SRC],0);
        xhandle(&hsrc,buf_src);
        if(hsrc.size!=0) {
            part_suffix_type = i;
            break;
        }
    }

    if(part_suffix_type == -1) {
        // printf("Can't find file:<%s>",args[ARGS_SRC]);
        return CP_FALSE;
    }

    printf("cpfile:%s\n", args[ARGS_SRC]);

    i = 0;
    // 计算分块总大小
    do {
        snprintf(buf_src,128,part_type[part_suffix_type],args[ARGS_SRC],i++);
        xhandle(&hsrc,buf_src);
        if(hsrc.size==0)
            break;
        total += hsrc.size;
    } while(part_suffix_type != 0); // 如果没有后缀则不继续

    i = 0;
    do {
        snprintf(buf_src,128,part_type[part_suffix_type],args[ARGS_SRC],i++);
        xhandle(&hsrc,buf_src);
        if(hsrc.size==0)
            break;
        
        do
        {
            actread = rw_handle_read(&hsrc,buffer,buffer_size);
            if(actread<=0){
                // printf("read ERROR:%lld\n",actread);
                break;
            }
            actwrite = rw_handle_write(&hdst,buffer,actread);
            if(actwrite<=0){
                printf("write ERROR:%lld\n",actread);
                break;
            }
            count += actwrite;
            printProgressBar(count,total);
        }while (actread==buffer_size);
    } while(part_suffix_type != 0); // 如果没有后缀则不继续

    putc('\n');

    if(count!=total)
        return CP_FALSE;
    return CP_TRUE;
}

U_BOOT_CMD(
	cpfile, 10, 1, do_cpfile,
	"cpfile in different device and different file system.",
	CMD_CPFILE_HELP
);
